!pip install numpy
!pip install matplotlib
!pip install nltk
!pip install dwave-ocean-sdk
!pip install openjij
!pip install pyqubo

import numpy as np
import matplotlib.pyplot as plt
import nltk
import pandas as pd
from nltk.corpus import wordnet as wn
from gensim.models import KeyedVectors
from pyqubo import Array, Constraint, Placeholder, Binary

words = [
    "art",
    "change_ringing",
    "peculiar",
    "english",
    "most",
    "english",
    "peculiarity",
    "unintelligible",
    "rest",
    "world",
]

correct_keys = {
    0: ["art%1:09:00::"],
    1: ["change_ringing%1:04:00::"],
    2: ["peculiar%5:00:00:specific:00", "peculiar%5:00:00:characteristic:00"],
    3: ["english%1:18:00::"],
    4: ["most%3:00:02::"],
    5: ["english%3:01:00::"],
    6: ["peculiarity%1:07:02::", "peculiarity%1:09:00::"],
    7: ["unintelligible%5:00:00:incomprehensible:00"],
    8: ["rest%1:24:00::"],
    9: ["world%1:14:02::"],
}

# WordNetのダウンロード
nltk.download("wordnet")
nltk.download("averaged_perceptron_tagger")
nltk.download("punkt")
nltk.download("omw-1.4")

word_data = {}
for i, word in enumerate(words):
    word_data[i] = {"word": word}

sense_data = {}
for i, word in enumerate(words):
    # 定義文の追加
    for k, synset in enumerate(wn.synsets(word)):
        sense_data[(i, k)] = {"sense": synset.definition()}

    word_data[i]["num_sense"] = k + 1

    # 語義キーの追加
    for k, synset in enumerate(wn.lemmas(word)):
        if synset.key() in correct_keys[i]:
            sense_data[(i, k)]["key"] = synset.key()
            sense_data[(i, k)]["correct"] = "*"
        else:
            sense_data[(i, k)]["key"] = synset.key()
            sense_data[(i, k)]["correct"] = ""

df = pd.DataFrame.from_dict(sense_data, orient="index")

df.head()

# 学習モデルのダウンロード(時間がかかります)
!wget http://wikipedia2vec.s3.amazonaws.com/models/en/2018-04-20/enwiki_20180420_100d.txt.bz2
# ファイルの解凍
!bzip2 -d enwiki_20180420_100d.txt.bz2

# モデルの読み込み(時間がかかります)
wikimodel = KeyedVectors.load_word2vec_format("enwiki_20180420_100d.txt", binary=False)

wikimodel["art"][:10]

# ベクトルの合成
wikimodel["change_ringing"] = wikimodel["change"] + wikimodel["ringing"]

for i, word in enumerate(words):
    word_data[i]["vector"] = wikimodel[word]

# cos類似度を計算
def cos_similarity(vec1, vec2):
    result = np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
    return result

# 語義ベクトルの作成
def sense2vec(parts, delta, w_data=word_data, s_data=sense_data):
    """
    parts: 内容語とする品詞
    delta: 閾値
    """
    for i, k in s_data.keys():
        morph = nltk.word_tokenize(s_data[i, k]["sense"])  # 文を単語に分割
        pos = nltk.pos_tag(morph)  # 単語の品詞を取得

        sum_vec = count = 0
        for p in pos:
            part_of_word = p[0]  # 内容語
            part_of_speech = p[1][0]  # 内容語の品詞

            if (
                (part_of_speech in parts)
                and (part_of_word in wikimodel)
                and (
                    cos_similarity(w_data[i]["vector"], wikimodel[part_of_word]) > delta
                )
            ):
                sum_vec += wikimodel[p[0]]
                count += 1

        sense_vec = sum_vec / count
        s_data[i, k]["vector"] = sense_vec

    return s_data

parts = ["J", "N", "R", "V"]
sense_data = sense2vec(parts=parts, delta=0)

def create_c(w_data=word_data, s_data=sense_data):
    c = {}
    for i, k in s_data.keys():
        c[i, k] = cos_similarity(w_data[i]["vector"], s_data[i, k]["vector"])
        s_data[i, k]["cos"] = round(c[i, k], 2)

    return c, s_data

cos, sense_data = create_c()

df = pd.DataFrame.from_dict(
    sense_data, orient="index", columns=["sense", "correct", "cos"]
)

df.head()

def create_b(w_data=word_data, s_data=sense_data):
    b = {
        (i, j, n): cos_similarity(w_data[i]["vector"], s_data[j, n]["vector"])
        for i in w_data.keys()
        for j, n in s_data.keys()
    }
    return b

def create_h(w_data=word_data, s_data=sense_data):
    h = {
        (i, m, j, n): cos_similarity(s_data[i, m]["vector"], s_data[j, n]["vector"])
        for i, m in s_data.keys()
        for j, n in s_data.keys()
    }
    return h

def create_e(w_data=word_data, s_data=sense_data):
    e = {
        (i, j): cos_similarity(w_data[i]["vector"], w_data[j]["vector"])
        for i in w_data.keys()
        for j in w_data.keys()
    }
    return e

def create_r(lam1, lam2, lam3, c=cos):
    b = create_b()
    h = create_h()
    e = create_e()

    # 値が大きくなりすぎないように正規化を行っている
    norm = lam1 * 2 + lam2 * 1 + lam3 * 3  # rの最大値

    r = {
        (i, m, j, n): (
            lam1 * (b[i, j, n] + b[j, i, m])
            + lam2 * h[i, m, j, n]
            + lam3 * (c[i, m] + e[i, j] + c[j, n])
        )
        / norm
        for i, m, j, n in h.keys()
    }
    return r

rel = create_r(lam1=1.0, lam2=1.0, lam3=1.0)

def formulation(w_data=word_data, s_data=sense_data, c=cos, r=rel):
    # バイナリ変数の作成
    x = {(i, k): Binary(f"x[{i}][{k}]") for i, k in s_data.keys()}

    # 目的関数の第1項
    H1 = np.sum([c[i, k] * x[i, k] for i, k in c.keys()])

    # 目的関数の第2項
    H2 = np.sum([r[i, u, j, v] * x[i, u] * x[j, v] for i, u, j, v in r.keys()])

    # 目的関数の第3項
    H3 = np.sum(
        [
            (np.sum([x[i, k] for k in range(w_data[i]["num_sense"])]) - 1) ** 2
            for i in w_data.keys()
        ]
    )

    # 最小化したい目的関数
    H = (
        -Placeholder("alpha") * H1
        - Placeholder("beta") * H2
        + Placeholder("gamma") * Constraint(H3, "H3")
    )
    model = H.compile()

    return model

# SAの場合
# from neal import SimulatedAnnealingSampler
# sampler = SimulatedAnnealingSampler()

# D-Waveマシンの場合
# from dwave.system import DWaveCliqueSampler
# sampler_config = {'solver': 'Advantage_system4.1', 'token': 'YOUR_TOKEN'}
# sampler = DWaveCliqueSampler(**sampler_config)

# Opnejijの場合
import openjij as oj
sampler = oj.SQASampler()

def extract_feasible_samples(decoded_samples, print_broken=False):
    feasible_samples = []
    for sample in decoded_samples:
        constraints = sample.constraints(only_broken=True)

        if len(constraints) == 0:
            feasible_samples.append(sample)
        elif print_broken:
            print(constraints)

    return feasible_samples

def optimization(alpha, beta, gamma, num_reads):
    model = formulation()

    feed_dict = {"alpha": alpha, "beta": beta, "gamma": gamma}  # 制約項の係数
    qubo, offset = model.to_qubo(feed_dict=feed_dict)

    sampleset = sampler.sample_qubo(qubo, num_reads=num_reads)

    decoded_samples = model.decode_sampleset(sampleset.aggregate(), feed_dict)
    feasible_samples = extract_feasible_samples(decoded_samples)

    return feasible_samples

feasible_samples = optimization(alpha=1.0, beta=1.0, gamma=30.0, num_reads=1000)
print("実行可能解:", len(feasible_samples))

def output_prob(feasible_samples, w_data=word_data, s_data=sense_data):
    # 確率の計算
    for i, k in s_data.keys():
        s_data[i, k]["prob(%)"] = 0

        for sample_info in feasible_samples:
            if sample_info.sample[f"x[{i}][{k}]"] == 1:
                s_data[i, k]["prob(%)"] += 1 / len(feasible_samples) * 100

    # 確率が1番高い語義を記録
    for i in w_data.keys():
        max_prob = 0

        for k in range(w_data[i]["num_sense"]):
            s_data[i, k]["prob(%)"] = round(s_data[i, k]["prob(%)"], 1)
            s_data[i, k]["max"] = ""

            if max_prob < s_data[i, k]["prob(%)"]:
                max_index = (i, k)
                max_prob = s_data[i, k]["prob(%)"]

        s_data[max_index]["max"] = "*"

    return s_data

sense_data = output_prob(feasible_samples=feasible_samples)
df = pd.DataFrame.from_dict(
    sense_data, orient="index", columns=["sense", "correct", "prob(%)", "max", "cos"]
)

df.head()

def count_correct(s_data):
    count = 0
    for i, k in s_data.keys():
        if s_data[i, k]["correct"] == "*" and s_data[i, k]["max"] == "*":
            count += 1
    return count

print("単語の数:", len(word_data))
print("正解数:", count_correct(s_data=sense_data))

def count_correct_per_sample(
    lam1, lam2, lam3, alpha, beta, gamma, num_reads=1000, s_data=sense_data
):
    """
    6個のハイパーパラメータを入力とし、1サンプル当たりの正解数を返す関数
    """
    r = create_r(lam1=lam1, lam2=lam2, lam3=lam3)
    feasible_samples = optimization(
        alpha=alpha, beta=beta, gamma=gamma, num_reads=num_reads
    )

    if len(feasible_samples) == 0:
        return 0

    count_list = []
    for sample_info in feasible_samples:
        count = 0
        for i, k in s_data.keys():
            if (sample_info.sample[f"x[{i}][{k}]"] == 1) and (
                s_data[i, k]["correct"] == "*"
            ):
                count += 1

        count_list.append(count)

    return np.average(count_list)

def objective(trial):
    param = {
        # rのハイパーパラメータ
        "lam1": trial.suggest_float("lam1", 0, 30, step=0.1),
        "lam2": trial.suggest_float("lam2", 0, 30, step=0.1),
        "lam3": trial.suggest_float("lam3", 0, 30, step=0.1),
        # 目的関数のハイパーパラメータ
        "alpha": trial.suggest_float("alpha", 0, 30, step=0.1),
        "beta": trial.suggest_float("beta", 0, 30, step=0.1),
        "gamma": trial.suggest_int("gamma", 30, 100),
    }

    return count_correct_per_sample(
        lam1=param["lam1"],
        lam2=param["lam2"],
        lam3=param["lam3"],
        alpha=param["alpha"],
        beta=param["beta"],
        gamma=param["gamma"],
    )

!pip install optuna

import optuna

study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=50)

def show_graph(study):
    epoches = []
    values = []
    best = 0
    for i in study.trials:
        if best < i.value:
            best = i.value
        epoches.append(i.number + 1)
        values.append(best)

    plt.plot(epoches, values, color="red")
    plt.title("optuna")
    plt.xlabel("trial")
    plt.ylabel("num_correct / sample")
    plt.show()

show_graph(study)

print(study.best_params)

rel = create_r(
    lam1=study.best_params["lam1"],
    lam2=study.best_params["lam2"],
    lam3=study.best_params["lam3"],
)

feasible_samples = optimization(
    alpha=study.best_params["alpha"],
    beta=study.best_params["beta"],
    gamma=study.best_params["gamma"],
    num_reads=1000,
)
print("実行可能解の数:", len(feasible_samples))

sense_data = output_prob(feasible_samples=feasible_samples)
df = pd.DataFrame.from_dict(
    sense_data, orient="index", columns=["sense", "correct", "prob(%)", "max", "cos"]
)

df.head()

print("単語の数:", len(word_data))
print("正解数:", count_correct(s_data=sense_data))
print("")

import gurobipy as gu

def use_gurobi(lam1, lam2, lam3, alpha, beta, gamma, s_data, c=cos, w_data=word_data):
    Model = gu.Model("wsd")
    Model.setParam("OutputFlag", 0)  # ログを非表示

    # 変数の定義
    x = {}
    for i, k in s_data.keys():
        x[i, k] = Model.addVar(vtype="B", name=f"x[{i},{k}]")
    Model.update()

    # 目的関数の定義
    r = create_r(lam1=lam1, lam2=lam2, lam3=lam3)
    H1 = gu.quicksum(c[i, k] * x[i, k] for i, k in c.keys())
    H2 = gu.quicksum(r[i, u, j, v] * x[i, u] * x[j, v] for i, u, j, v in r.keys())
    H3 = gu.quicksum(
        (gu.quicksum(x[i, k] for k in range(w_data[i]["num_sense"])) - 1) ** 2
        for i in w_data.keys()
    )
    Model.setObjective(-alpha * H1 - beta * H2 + gamma * H3, gu.GRB.MINIMIZE)
    Model.update()

    Model.optimize()

    # 記録
    for i, k in s_data.keys():
        val = x[i, k].X
        if val == 1:
            s_data[i, k]["ans"] = "*"
        else:
            s_data[i, k]["ans"] = ""
    return s_data

# 最適なパラメータで厳密解を求める
sense_data = use_gurobi(
    lam1=study.best_params["lam1"],
    lam2=study.best_params["lam2"],
    lam3=study.best_params["lam3"],
    alpha=study.best_params["alpha"],
    beta=study.best_params["beta"],
    gamma=study.best_params["gamma"],
    s_data=sense_data,
)

df = pd.DataFrame.from_dict(
    sense_data, orient="index", columns=["sense", "correct", "ans"]
)

df.head()

print("単語の数:", len(word_data))
print("正解数:", count_correct(s_data=sense_data))